package com.demo.图;

import com.demo.其它.UnionFind;
import com.demo.图.model.Edge;
import com.demo.图.model.Graph;
import com.demo.图.model.Vertex;

import java.util.*;

/**
 * 图的最小生成树
 * @author cl
 */
public class GraphMST {

    /** 使用 Kruskal 算法 返回最小生成树 */
    public static <T> Graph<T> Kruskal(Graph<T> graph) {
        Objects.requireNonNull(graph);
        Graph<T> result = new Graph<>();
        // 需要使用并查集 并查集数据映射 做成
        // 前驱节点的做成
        int i = 0;
        int[] pre = new int[graph.vertexMap.size()];
        int[] rank = new int[graph.vertexMap.size()];
        Map<T, Integer> tmpMapNode = new HashMap<>();
        for (Map.Entry<T, Vertex<T>> entry : graph.vertexMap.entrySet()) {
            // 1. 节点值标记转换为数字 便于处理
            tmpMapNode.put(entry.getKey(), i);
            // 默认它的前驱节点是它自己
            rank[i] = 1;
            pre[i] = i;
            i++;
        }

        // 创建并查集
        UnionFind unionFind = new UnionFind(pre, rank);
        // 将边存储堆中
        PriorityQueue<Edge<T>> edgeHeap = new PriorityQueue<>(Comparator.comparingInt(e -> e.weight));
        edgeHeap.addAll(graph.edges);
        while (!edgeHeap.isEmpty()) {
            // 取得最小的边
            Edge<T> tmpEdge = edgeHeap.poll();
            // 查看边的from 和 to 是不是一个并查集 是的话就跳过，不是的话合并 加入
            if (!unionFind.isSame(tmpMapNode.get(tmpEdge.from.data), tmpMapNode.get(tmpEdge.to.data))) {
                unionFind.union(tmpMapNode.get(tmpEdge.from.data), tmpMapNode.get(tmpEdge.to.data));
                result.addEdge(tmpEdge);
            }
        }
        return result;
    }

    /** 使用 Prim 算法 返回最小生成树 */
    /** 1）以某一个点开始，寻找当前该点可以访问的所有的边；
     *  2）在已经寻找的边中发现最小边，这个边必须有一个点还没有访问过，将还没有访问的点加入我们的集合，记录添加的边；
     *  3）寻找当前集合可以访问的所有边，重复2的过程，直到没有新的点可以加入；
     *  4）此时由所有边构成的树即为最小生成树。
     */
    public static <T> Graph<T> prim(Graph<T> graph) {
        Objects.requireNonNull(graph);
        Graph<T> result = new Graph<>();
        HashSet<Vertex<T>> selected = new HashSet<>();
        HashSet<Edge<T>> allEdge = new HashSet<>();
        for(Vertex<T> node : graph.vertexMap.values()) {
            if (!selected.contains(node)) {
                selected.add(node);
                allEdge.addAll(node.edges);
                while (true) {
                    Edge<T> minEdge = null;
                    // 找到一个最小的边
                    for (Edge<T> tmpEdge : allEdge) {
                        if (!selected.contains(tmpEdge.to)) {
                            if (minEdge == null) {
                                minEdge = tmpEdge;
                            } else if (minEdge.weight > tmpEdge.weight) {
                                minEdge = tmpEdge;
                            }
                        }
                    }
                    if (minEdge != null) {
                        // 最小边的节点入选
                        selected.add(minEdge.to);
                        // 最小边添加到结果集里面
                        result.addEdge(minEdge);
                        // 记录最小边节点的所有边
                        allEdge.addAll(minEdge.to.edges);
                    } else {
                        // 找不到最小边了 也就是没有了
                        break;
                    }
                }
            }
        }
        return result;
    }
}
